// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <vector>
#include <pollux/core/expressions.h>
#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/function_metadata.h>
#include <pollux/expression/function_signature.h>
#include <pollux/vector/selectivity_vector.h>
#include <pollux/vector/simple_vector.h>

#include <melon/synchronized.h>

namespace kumo::pollux::exec {
    class Expr;
    class ExprSet;
    class EvalCtx;

    // Superclass for functions which are written to run on whole vectors.
    class VectorFunction {
    public:
        virtual ~VectorFunction() = default;

        /// Evaluate the function on the specified rows.
        ///
        /// The rows may or may not be contiguous.
        ///
        /// Single-argument deterministic functions may receive their only argument
        /// vector as flat or constant, but not dictionary encoded.
        ///
        /// Single-argument functions that specify null-in-null-out behavior, e.g.
        /// defaultNullBehavior is true, will never see a null row in 'rows'.
        /// Hence, they can safely assume that args[0] vector is flat or constant and
        /// has no nulls in specified positions.
        ///
        /// Multi-argument functions that specify null-in-null-out behavior will never
        /// see a null row in any of the arguments. They can safely assume that there
        /// are no nulls in any of the arguments in specified positions.
        ///
        /// If context.isFinalSelection() is false, the result may have been
        /// partially populated for the positions not specified in rows. The function
        /// must take care not to overwrite these values. This happens when evaluating
        /// conditional expressions. Consider if(a = 1, f(b), g(c)) expression. The
        /// engine first evaluates a = 1, then invokes f(b) for passing rows, then
        /// invokes g(c) for the rest of the rows. f(b) is invoked with null result,
        /// but g(c) is called with partially populated result which has f(b) values
        /// in some positions. Function g must preserve these values.
        ///
        /// Use context.isFinalSelection() to determine whether partially populated
        /// results must be preserved or not.
        ///
        /// If 'result' is not null, it can be dictionary, constant or sequence
        /// encoded and therefore may be read-only. Call BaseVector::ensure_writable
        /// before writing into the "result", e.g.
        ///
        ///    BaseVector::ensure_writable(&rows, caller->type(), context.pool(),
        ///        result);
        ///
        /// If 'result' is null, the function can allocate a new vector using
        /// BaseVector::create or reuse memory of its argument(s). Memory reuse is
        /// safe if the function returns its argument unmodified or the argument and
        /// its contents are uniquely referenced.
        virtual void apply(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args, // Not using const ref so we can reuse args
            const TypePtr &outputType,
            EvalCtx &context,
            VectorPtr &result) const = 0;

        /// Returns true if (1) supports evaluation on all constant inputs of size >
        /// 1; (2) returns flat or constant result when inputs are all flat, all
        /// constant or a mix of flat and constant; (3) guarantees that if all inputs
        /// are not null, the result is also not null.
        virtual bool supportsFlatNoNullsFastPath() const {
            return false;
        }

        // The evaluation engine will scan and set the string encoding of the
        // specified input arguments when presented if their type is VARCHAR before
        // applying the function
        virtual std::vector<size_t> ensureStringEncodingSetAt() const {
            return {};
        }

        // Same as ensureStringEncodingSetAt but for all string inputs
        virtual bool ensureStringEncodingSetAtAllInputs() const {
            return false;
        }

        // If set, the string encoding of the results will be set by propagating
        // the string encoding from all provided string inputs.
        virtual bool propagateStringEncodingFromAllInputs() const {
            return false;
        }

        // If set, the string encoding of the results will be set by propagating
        // the specified inputs string encodings if presented.
        // If one of the specified inputs have its encoding not determined, the
        // encoding of the result is not determined.
        virtual std::optional<std::vector<size_t> > propagateStringEncodingFrom()
        const {
            return std::nullopt;
        }

        virtual FunctionCanonicalName getCanonicalName() const {
            return FunctionCanonicalName::kUnknown;
        }
    };

    /// Vector function that generates the specified error for every row. Use this
    /// to hold an error generated while processing inputs for a stateful vector
    /// function. Such errors should not be reported until VectorFunction::apply is
    /// called to avoid signaling errors for expressions which end up not needing to
    /// evaluate the function.
    class AlwaysFailingVectorFunction final : public VectorFunction {
    public:
        explicit AlwaysFailingVectorFunction(std::exception_ptr exceptionPtr)
            : exceptionPtr_{std::move(exceptionPtr)} {
        }

        void apply(
            const SelectivityVector &rows,
            std::vector<VectorPtr> & /*args*/,
            const TypePtr & /* outputType */,
            EvalCtx &context,
            VectorPtr & /*resultRef*/) const final {
            context.setErrors(rows, exceptionPtr_);
        }

    private:
        std::exception_ptr exceptionPtr_;
    };

    // This functions is used when we know a function will never be called because
    // it is default null with a literal null input value. For example like(c0,
    // null).
    class ApplyNeverCalled final : public VectorFunction {
        void apply(
            const SelectivityVector &,
            std::vector<VectorPtr> &,
            const TypePtr &,
            EvalCtx &,
            VectorPtr &) const final {
            POLLUX_UNREACHABLE("Not expected to be called.");
        }
    };

    // Factory for functions which are template generated from simple functions.
    class SimpleFunctionAdapterFactory {
    public:
        virtual std::unique_ptr<VectorFunction> createVectorFunction(
            const std::vector<TypePtr> &inputTypes,
            const std::vector<VectorPtr> &constantInputs,
            const core::QueryConfig &config) const = 0;

        virtual ~SimpleFunctionAdapterFactory() = default;
    };

    /// Returns the function metadata with the specified name. Returns std::nullopt
    /// if there is no function with the specified name.
    std::optional<VectorFunctionMetadata> getVectorFunctionMetadata(
        const std::string &name);

    /// Returns a list of signatures supported by VectorFunction with the specified
    /// name. Returns std::nullopt if there is no function with the specified name.
    std::optional<std::vector<FunctionSignaturePtr> > getVectorFunctionSignatures(
        const std::string &name);

    /// Given name of vector function and argument types, returns
    /// the return type if function exists and have a signature that binds to the
    /// input types otherwise returns nullptr.
    TypePtr resolveVectorFunction(
        const std::string &functionName,
        const std::vector<TypePtr> &argTypes);

    std::optional<std::pair<TypePtr, VectorFunctionMetadata> >
    resolveVectorFunctionWithMetadata(
        const std::string &functionName,
        const std::vector<TypePtr> &argTypes);

    /// Returns an instance of VectorFunction for the given name, input types and
    /// optionally constant input values.
    /// constantInputs should be empty if there are no constant inputs.
    /// constantInputs should be aligned with inputTypes if there is at least one
    /// constant input; non-constant inputs should be represented as nullptr;
    /// constant inputs must be instances of ConstantVector.
    std::shared_ptr<VectorFunction> getVectorFunction(
        const std::string &name,
        const std::vector<TypePtr> &inputTypes,
        const std::vector<VectorPtr> &constantInputs,
        const core::QueryConfig &config);

    std::optional<
        std::pair<std::shared_ptr<VectorFunction>, VectorFunctionMetadata> >
    getVectorFunctionWithMetadata(
        const std::string &name,
        const std::vector<TypePtr> &inputTypes,
        const std::vector<VectorPtr> &constantInputs,
        const core::QueryConfig &config);

    /// Registers stateless VectorFunction. The same instance will be used for all
    /// expressions.
    /// Returns true iff an new function is inserted
    bool registerVectorFunction(
        const std::string &name,
        std::vector<FunctionSignaturePtr> signatures,
        std::unique_ptr<VectorFunction> func,
        VectorFunctionMetadata metadata = {},
        bool overwrite = true);

    // Represents arguments for stateful vector functions. Stores element type, and
    // the constant value (if supplied).
    struct VectorFunctionArg {
        const TypePtr type;
        const VectorPtr constantValue;
    };

    using VectorFunctionFactory = std::function<std::shared_ptr<VectorFunction>(
        const std::string &name,
        const std::vector<VectorFunctionArg> &inputArgs,
        const core::QueryConfig &config)>;

    struct VectorFunctionEntry {
        std::vector<FunctionSignaturePtr> signatures;
        VectorFunctionFactory factory;
        VectorFunctionMetadata metadata;
    };

    // TODO: Use melon::Singleton here
    using VectorFunctionMap =
    melon::Synchronized<melon::F14FastMap<std::string, VectorFunctionEntry> >;

    VectorFunctionMap &vectorFunctionFactories();

    // A template to simplify making VectorFunctionFactory for a function that has a
    // constructor that takes inputTypes and constantInputs
    //
    // Example of usage:
    //
    // POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION(
    //    udf_my_function, exec::makeVectorFunctionFactory<MyFunction>());
    template<typename T>
    VectorFunctionFactory makeVectorFunctionFactory() {
        return [](const std::string &name,
                  const std::vector<VectorFunctionArg> &inputArgs,
                  const core::QueryConfig & /*config*/) {
            return std::make_shared<T>(name, inputArgs);
        };
    }

    // Registers stateful VectorFunction. New instance is created for each
    // expression using input types and optionally constant values for some inputs.
    // When overwrite is true, any previously registered VectorFunction with the
    // name is replaced.
    // Returns true iff the function was inserted
    bool registerStatefulVectorFunction(
        const std::string &name,
        std::vector<FunctionSignaturePtr> signatures,
        VectorFunctionFactory factory,
        VectorFunctionMetadata metadata = {},
        bool overwrite = true);

    /// An expression re-writer that takes an expression and returns an equivalent
    /// expression or nullptr if re-write is not possible.
    using ExpressionRewrite = std::function<core::TypedExprPtr(core::TypedExprPtr)>;

    /// Returns a list of registered re-writes.
    std::vector<ExpressionRewrite> &expressionRewrites();

    /// Appends a 'rewrite' to 'expressionRewrites'.
    ///
    /// The logic that applies re-writes is very simple and assumes that all
    /// rewrites are independent. Re-writes are applied to all expressions starting
    /// at the root and going down the hierarchy. For each expression, rewrites are
    /// applied in the order they were registered. The first rewrite that returns
    /// non-null result terminates the re-write for this particular expression.
    void registerExpressionRewrite(ExpressionRewrite rewrite);
} // namespace kumo::pollux::exec

// Private. Return the external function name given a UDF tag.
#define _POLLUX_REGISTER_FUNC_NAME(tag) registerVectorFunction_##tag

// Declares a vectorized UDF function given a tag. Goes into the UDF .cpp file.
#define POLLUX_DECLARE_VECTOR_FUNCTION(tag, signatures, function) \
  void _POLLUX_REGISTER_FUNC_NAME(tag)(const std::string& name) { \
    kumo::pollux::exec::registerVectorFunction(               \
        (name), (signatures), (function));                       \
  }

#define POLLUX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(             \
    tag, signatures, metadata, function)                         \
  void _POLLUX_REGISTER_FUNC_NAME(tag)(const std::string& name) { \
    kumo::pollux::exec::registerVectorFunction(               \
        (name), (signatures), (function), (metadata));           \
  }

// Declares a stateful vectorized UDF.
#define POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION(tag, signatures, function) \
  void _POLLUX_REGISTER_FUNC_NAME(tag)(const std::string& name) {          \
    kumo::pollux::exec::registerStatefulVectorFunction(                \
        (name), (signatures), (function));                                \
  }

#define POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION_WITH_METADATA(    \
    tag, signatures, metadata, function)                         \
  void _POLLUX_REGISTER_FUNC_NAME(tag)(const std::string& name) { \
    kumo::pollux::exec::registerStatefulVectorFunction(       \
        (name), (signatures), (function), (metadata));           \
  }

// Registers a vectorized UDF associated with a given tag.
// This should be used in the same namespace the declare macro is used in.
#define POLLUX_REGISTER_VECTOR_FUNCTION(tag, name)                   \
  {                                                                 \
    extern void _POLLUX_REGISTER_FUNC_NAME(tag)(const std::string&); \
    _POLLUX_REGISTER_FUNC_NAME(tag)(name);                           \
  }
